#! /usr/bin/python3
# -*- coding: UTF-8 -*-
"""
  作者: 小肥爬爬
  简书: https://www.jianshu.com/u/db796a501972
  gitee: https://gitee.com/xiaofeipapa/python-toolkit
  您可以自由转载此博客文章, 恳请保留原链接, 谢谢!
"""
import os, os.path
import fab_utils
from fab_utils import FabException
from fabric import Connection


def hint(msg):

    print('---- ' + msg + '\n')


def title(msg):

    print('-----------------------------------------------------------------------------')
    print(msg)
    print('-----------------------------------------------------------------------------\n')


class UserTask:
    """
    在远程服务器增加用户
    """

    def __init__(self, user, host):

        self.conn = Connection(host, user=user)
        self.host = host

        self._prepare()

    def _prepare(self):

        # 确保有 pwgen 依赖
        cmd = 'apt-get install pwgen'
        fab_utils.run(self.conn, cmd)
        hint('安装pwgen完成')

    def add_user(self, username):
        """
        增加用户, 并设置密码.
        :param username: 用户名
        :return: 登录密码
        """

        if self._is_user_exist(username):
            title('用户 ' + username + ' 已经存在')
            fab_utils.close_connection(self.conn)
            return None

        # 生成10位密码
        cmd = 'pwgen -s 10 1'
        password = fab_utils.run(self.conn, cmd)

        # 生成用户. 该命令只在centos 测试过
        hint('正在新增用户 ' + username + '......')
        fab_utils.ubuntu_adduser(self.conn, username, password)
        hint('已增加用户 ' + username)

        # 退出root登录
        fab_utils.close_connection(self.conn)

        # 新用户登录进行测试
        self.verify_login(username, self.host, password)

        return password

    # 验证该用户是否可以登录
    def verify_login(self, user, host, password):
        user_conn = Connection(user=user, host=host, connect_kwargs={'password': password})
        result = fab_utils.run(user_conn, 'pwd')
        if result == '/home/' + user:
            hint('用户 ' + user + '登录 ' + host + ' 成功')
        else:
            hint('用户 ' + user + '登录 ' + host + ' 失败!!!!!')

        fab_utils.close_connection(user_conn)

    def _is_user_exist(self, username):

        # 如果存在会返回 uid=1000(test) gid=1000(test) groups=1000(test) 之类
        # 不存在则抛出异常
        cmd = 'id ' + username
        try:
            fab_utils.run(self.conn, cmd)
            return True
        except FabException:
            return False


class RemoteAutoTask:
    """
    在远程服务器运行自动化任务
    """

    def __init__(self, user, host, task, rerun):
        """
        :param task: 任务名称
        :param from_dir: 该任务的主目录
        """
        self.task = task
        self.conn = Connection(host, user=user)

        self._prepare(rerun)

        title('准备在 ' + host + ' 进行任务: ' + task + ' 主目录: ' + self.task_home)

    def _prepare(self, rerun):
        """
        检查任务是否已经启动过
        :return:
        """
        # 必须要先run连接远程服务器, 否则 conn.client.exec_command 会报错
        fab_utils.run(self.conn, 'ls /tmp')

        # 生成 /root/data 目录
        base_dir = '/root/data'
        fab_utils.check_has_dir(self.conn, base_dir)

        # 检查该task 是否已经执行过. 如果已经执行过, 在 /root/data/task 目录下有 done.txt 文件
        task_home = base_dir + '/' + self.task
        fab_utils.check_has_dir(self.conn, task_home)
        self.task_home = task_home

        task_file = task_home + '/done.txt'
        if not fab_utils.check_has_file(self.conn, task_file):
            return

        # 如果存在, 又不指定重新运行, 则退出
        if not rerun:
            title('此程序已经在目标机器运行过了. 如果想再次运行, 请设置启动参数 -r true')
            exit(0)

    def finish(self):

        # 在主目录下生成 done.txt
        fab_utils.run(self.conn, 'touch ' + self.task_home + '/done.txt')

        fab_utils.close_connection(self.conn)

        print('任务 ' + self.task + ' 已完成!! ')


class SetupJava(RemoteAutoTask):
    def __init__(self, user, host, rerun):

        task = 'install_jdk'

        RemoteAutoTask.__init__(self, user, host, task, rerun)

    # 安装 prometheus jvm exporter
    def install_jvm_exporter(self, exporter_file, exporter_home):

        with self.conn.cd(self.task_home):

            # 检查下载
            file_path = self.task_home + '/' + exporter_file
            file_md5 = 'ac475ee988c8a52d2310f073e67aca61'
            down_link = 'https://repo1.maven.org/maven2/io/prometheus/jmx/jmx_prometheus_javaagent/0.15.0/'\
                     + exporter_file
            fab_utils.check_download(self.conn, file_path, file_md5, down_link)

            # 复制
            self.conn.run('/bin/cp -rf ' + exporter_file + ' ' + exporter_home)

        return os.path.join(exporter_home, exporter_file)

    # 结合两个任务
    def run_all(self, exporter_file, exporter_home, jdk_file, jvm_dir):

        # 检查jvm 目录
        fab_utils.check_has_dir(self.conn, jvm_dir)

        # 下载 jvm exporter jar
        exporter_jar = self.install_jvm_exporter(exporter_file, exporter_home)

        # 安装jdk
        self.install_jdk(jdk_file, jvm_dir)
        #
        # # 写入环境变量
        self._check_write_env_var(self.conn, jvm_dir, exporter_jar)

    def install_jdk(self, jdk_file, jvm_dir):

        with self.conn.cd(self.task_home):

            # 检查下载jdk
            file_path = self.task_home + '/' + jdk_file
            file_md5 = 'ef599e322eee42f6769991dd3e3b1a31'
            down_link = 'https://repo.huaweicloud.com/java/jdk/8u181-b13/' + jdk_file
            fab_utils.check_download(self.conn, file_path, file_md5, down_link)

            # 创建jvm目录
            fab_utils.check_has_dir(self.conn, jvm_dir)

            # 解压. 华为的这个包解压后名字不同, 所以手动指定一个目录
            hint('正在解压 jdk ......')
            fab_utils.run(self.conn, 'tar xvf ' + jdk_file + ' -C ' + jvm_dir + '  --strip-components 1 ')

    def _check_write_env_var(self, conn, jvm_dir, exporter_jar_path):

        """
        一开始把环境变量写到 /etc/profile 不成功, 原因是 /etc/profile 在用户第一次登录的时候才读取.
         对于新增的变量, 应该写到 /etc/bashrc
        """
        # system_file = '/etc/profile'
        system_file = '/etc/bashrc'

        run_result = conn.run('echo $JAVA_HOME ', hide=True).stdout.strip()
        if run_result == jvm_dir:
            hint(system_file + ' 已经包含 JAVA_HOME, 不需要重新设置')
        else:
            # 追加内容到系统配置
            conn.run("echo '# ===== java 配置' >> " + system_file)
            conn.run("echo 'export JAVA_HOME=" + jvm_dir + "' >> " + system_file)

            # 将 exporter jar 也加入到 classpath 里
            export = 'export CLASSPATH=.:$JAVA_HOME/lib/dt.jar:$JAVA_HOME/lib/tools.jar:$JAVA_HOME/jre/lib/rt.jar:'\
                     + exporter_jar_path
            conn.run("echo '" + export + "' >> " + system_file)
            conn.run("echo 'export PATH=$PATH:$JAVA_HOME/bin' >> " + system_file)

            # 使环境变量生效
            conn.run('source ' + system_file)


def do_it(args):

    user = 'root'
    host = '192.168.145.24'

    jdk_file = 'jdk-8u181-linux-x64.tar.gz'
    jvm_dir = '/usr/local/jvm/jdk8'

    exporter_file = 'jmx_prometheus_javaagent-0.15.0.jar'
    exporter_home = '/usr/local/jvm/jmx_exporter'

    java_user = 'java_user'

    java_task = SetupJava(user, host, args.rerun)

    java_task.run_all(exporter_file, exporter_home, jdk_file, jvm_dir)

    # 增加java 用户
    user_task = UserTask(user, host)
    password = user_task.add_user(java_user)
    if password:
        title('安装完成!! java 用户: ' + java_user + '  登录密码: ' + password)


if __name__ == '__main__':

    import argparse

    parser = argparse.ArgumentParser('传入参数：***.py')
    parser.add_argument('-r', '--rerun', default=True)    # 开发时用
    # parser.add_argument('-r', '--rerun', default=False)

    args = parser.parse_args()

    rerun = args.rerun

    print('---- 参数 rerun: ', rerun)

    do_it(args)


